//Examples/reed_sol_rs_01.cpp

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "jerasure.h"
#include "reed_sol.h"

#define talloc(type, num) (type *) malloc(sizeof(type)*(num))

void usage(char *s)
{
	fprintf(stderr, "usage: reed_sol_rs_01 k m w - Does a simple Reed-Solomon coding example in GF(2^w).\n");
	fprintf(stderr, "       \n");
	fprintf(stderr, "       w must be 8, 16 or 32.  k+m must be <= 2^w.  It sets up a classic\n");
	fprintf(stderr, "       Vandermonde-based distribution matrix and encodes k devices of\n");
	fprintf(stderr, "       8 bytes each with it.  Then it decodes.\n");
	fprintf(stderr, "       \n");
	fprintf(stderr, "This demonstrates: RS_Classic_Generator()\n");
	fprintf(stderr, "                   JER_Matrix::Print()\n");
	fprintf(stderr, "                   JER_Slices::Encode()\n");
	fprintf(stderr, "                   JER_Slices::Decode()\n");
	if (s != NULL) fprintf(stderr, "%s\n", s);
	exit(1);
}

void print_data_and_coding(JER_Slices *slices)
{
	int i, j, x;
	int n, sp;
	int k, m, w, size;
	uint64_t l;
	JER_Gen_T * g;
	vector < string > states;
	int max_data_state_len; //max strlen of a data drive state
	int max_coding_state_len;
	int state_len;

	states.resize(3);
	states[0] = "up";
	states[1] = "down";
	states[2] = "unusable";

	max_data_state_len = max_coding_state_len = 2;
	for(i=0;i<slices->N;i++){
		state_len = strlen(states[slices->States[i]].c_str());
		if(i<slices->K){
			//data
			if(state_len>max_data_state_len){
				max_data_state_len = state_len;
			}
		}else{
			//coding
			if(state_len>max_coding_state_len){
				max_coding_state_len = state_len;
			}
		}
	}

	g = slices->G;
	if(g == NULL){
		return;
	}

	if(g->M == NULL){
		return;
	}

	k = slices->K;
	m = slices->N - k;
	w = g->M->W;
	size = slices->PacketSize * slices->PacketsPerSlice;

	if (k > m) n = k;
	else n = m;
	sp = size * 2 + size/(w/8) + 8+max_data_state_len+3;

	printf("%-*sCoding\n", sp, "Data");
	for (i = 0; i < n; i++) {
		if (i < k) {
			printf("D%-2d (%-*s):", i, max_data_state_len,states[slices->States[i]].c_str());
			for (j = 0; j < size; j+=(w/8)) {
				printf(" ");
				for (x = 0; x < w/8; x++){
					printf("%02x", (unsigned char)slices->Ptrs[i][j+x]);
				}
			}
			printf("    ");
		}
		else printf("%*s", sp, "");
		if (i < m) {
			printf("C%-2d (%-*s):", i, max_coding_state_len,states[slices->States[i+k]].c_str());
			for (j = 0 ; j < size; j+=(w/8)) {
				printf(" ");
				for (x = 0; x < w/8; x++){
					printf("%02x", (unsigned char)slices->Ptrs[i+k][j+x]);
				}
			}
		}
		printf("\n");
	}
	printf("\n");
}

int main(int argc, char **argv){
	uint64_t l;
	int k, w, i, j, m,ps;
	int size;
	JER_Matrix *jm;
	JER_Slices *slices;
	unsigned char *data;
	vector <int> erasures;
	JER_Matrix *dm;
	int *dm_ids;
	JER_Gen_T * g;

	if (argc != 4) usage(NULL);
	if (sscanf(argv[1], "%d", &k) == 0 || k <= 0) usage((char *)"Bad k");
	if (sscanf(argv[2], "%d", &m) == 0 || m <= 0) usage((char *)"Bad m");
	if (sscanf(argv[3], "%d", &w) == 0 || (w != 8 && w != 16 && w != 32)) usage((char *)"Bad w");
	if (w <= 16 && k + m > (1 << w)) usage((char *)"k + m is too big");
	ps = size = 8; 

	g = RS_Classic_Generator(k,m,w);

	printf("Last m rows of the Distribution Matrix:\n\n");
	g->M->Print();
	printf("\n");

	srand48(0);

	//allocate and create empty slices	 
	slices = new JER_Slices(k+m, k,ps,1,g);
	
	data = new unsigned char [size * (k + m)];

	for (i = 0; i < k; i++) {
		for(j = 0; j < size; j+=ps) {
			l = lrand48();
			l <<= 8*4;
			l += lrand48();

			memcpy(data + i*size + j, &l, ps);
		}
		slices->Ptrs[i] = data + i*size;
	}

	for (i = 0; i < m; i++) {
		slices->Ptrs[k + i] = data + k*size + i*size;
	}

	printf("Original data:\n\n");
	print_data_and_coding(slices);

	slices->Encode();

	printf("Encoding complete:\n\n");
	print_data_and_coding(slices);

	erasures.resize(m);
	l = 0;
	for (i = 0; i < m; ) {
		erasures[i] = lrand48()%(k+m);
		if (slices->States[erasures[i]] == 0) {
			slices->States[erasures[i]] = 1;
			memcpy(data + erasures[i]*size, &l, size);
			i++;
		}
	}

	printf("Erased %d random devices:\n\n", m);
	print_data_and_coding(slices);

	slices->Decode();

	printf("State of the system after decoding:\n\n");
	print_data_and_coding(slices);

	return 0;
}
